One of the important aspects of data analysis is making outputs that are easy to understand and aesthetically pleasing. In my work, I refer to this as making graphs “pretty”. Today I want to test whether LLMs have a concept of “prettiness”. I’m interested to see how they behave when I ask them to modify graph code to make it more aesthetically pleasing.
For reference, this is what the original graph looks like:
Show the code
import pandas as pdimport plotly_express as pxanx2010_df = pd.read_csv('data/global_anxiety_2010.csv')fig = px.choropleth(anx2010_df, locations="Code", color="Anxiety_percent", hover_name="Entity", color_continuous_scale=px.colors.sequential.Plasma)fig.show()
Test 1: GPT 4 Turbo
Show the code
from openai import OpenAIimport osimport pandas as pdfrom IPython.display import Markdown, displayfrom dotenv import load_dotenv# Load API key_ = load_dotenv()client = OpenAI()model ="gpt-4-turbo-2024-04-09"system_role ='''You are a consultant skilled at making visually appealing and easy to understandgraphs to present to clients.'''initial_prompt ='''Improve this code to make the graph better. Reply only with Python code.Here is the code to improve: import pandas as pdimport plotly_express as pxanx2010_df = pd.read_csv('data/global_anxiety_2010.csv')fig = px.choropleth(anx2010_df, locations="Code", color="Anxiety_percent", hover_name="Entity", color_continuous_scale=px.colors.sequential.Plasma)fig.show()'''messages = [ {"role": "system", "content": system_role} , {"role": "user", "content": initial_prompt},]completion = client.chat.completions.create( model=model, # the model to use messages=messages, # a list of dictionaries with the prompt n=3, # how many responses to generate temperature=0.7# temperature parameter)response = display(Markdown(f'-----\n#### Response {0}\n{resp.message.content}\n'))
Response 0
import pandas as pdimport plotly_express as px# Loading dataanx2010_df = pd.read_csv('data/global_anxiety_2010.csv')# Creating a choropleth map with improved featuresfig = px.choropleth(anx2010_df, locations="Code", color="Anxiety_percent", hover_name="Entity", title="Global Anxiety Rates in 2010", labels={"Anxiety_percent": "Anxiety Rate (%)"}, color_continuous_scale=px.colors.sequential.Plasma, projection="natural earth")# Enhancing layoutfig.update_layout( geo=dict( showframe=False, showcoastlines=True, projection_type='equirectangular' ), coloraxis_colorbar=dict( title="Anxiety Rate %", thicknessmode="pixels", thickness=15, lenmode="pixels", len=300 ))# Showing the plotfig.show()
Testing the code
Show the code
import pandas as pdimport plotly_express as px# Loading dataanx2010_df = pd.read_csv('data/global_anxiety_2010.csv')# Creating a choropleth map with improved featuresfig = px.choropleth(anx2010_df, locations="Code", color="Anxiety_percent", hover_name="Entity", title="Global Anxiety Rates in 2010", labels={"Anxiety_percent": "Anxiety Rate (%)"}, color_continuous_scale=px.colors.sequential.Plasma, projection="natural earth")# Enhancing layoutfig.update_layout( geo=dict( showframe=False, showcoastlines=True, projection_type='equirectangular' ), coloraxis_colorbar=dict( title="Anxiety Rate %", thicknessmode="pixels", thickness=15, lenmode="pixels", len=300 ))# Showing the plotfig.show()
Analysis
Here we can see that we have:
A title for the overall graph
A label for the anxiety rate legend
No more black box around the outside of the map
Test 2: Claude 3.5 Sonnet
Show the code
from dotenv import load_dotenvfrom anthropic import Anthropicimport base64import httpxfrom IPython.display import display, Image, Markdown# load api keys_ = load_dotenv()client = Anthropic()message = client.messages.create( model='claude-3-5-sonnet-20241022', max_tokens=1000, temperature=0.0, system='''You are a consultant skilled at making visually appealing and easy to understandgraphs to present to clients.''', messages=[ {"role": "user", "content": '''Improve this code to make the graph better. Reply only with Python code.Here is the code to improve: import pandas as pdimport plotly_express as pxanx2010_df = pd.read_csv('data/global_anxiety_2010.csv')fig = px.choropleth(anx2010_df, locations="Code", color="Anxiety_percent", hover_name="Entity", color_continuous_scale=px.colors.sequential.Plasma)fig.show()'''} ])response = display(Markdown(f'-----\n#### Response {0}\n{message.content[0].text}\n'))
Both models did pretty well with this task. A lot of a graph’s “prettiness” comes down to personal preference, but asking LLMs for various options could definitely help in the brainstorming process.